/* 
 * Copyright 2014 Josef Hertl.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package caetus.pkcs11.key;

import caetus.pkcs11.template.Template;
import sun.security.pkcs11.wrapper.CK_ATTRIBUTE;
import sun.security.pkcs11.wrapper.CK_MECHANISM;
import sun.security.pkcs11.wrapper.PKCS11;
import sun.security.pkcs11.wrapper.PKCS11Constants;
import sun.security.pkcs11.wrapper.PKCS11Exception;

/**
 * Represents a secret key.
 *
 * @author Josef Hertl
 */
public class SecretKey extends Key {

    /**
     * Only really supports block ciphers so far.
     */
    private int blockSize = 32;

    /**
     * @param sessionHandle The session handle in which all operations are
     * performed.
     * @param pkcs11 The pkcs object that performs all the operations.
     * @param cipherEnum The cipher of this key.
     */
    public SecretKey(long sessionHandle, PKCS11 pkcs11, CipherEnum cipherEnum) {
        super(sessionHandle, pkcs11, cipherEnum);
    }

    /**
     * Generates a new value for this key.
     *
     * @param template The template which is passed to the generation mechanism.
     * It is used to specify the attributes of the newly generated key.
     */
    public void generate(Template template) throws PKCS11Exception {
        if (template == null) {
            template = new Template();
            Template basicTemplate = templateMaker.getBasicSecretKeyTemplate(cipherEnum);
            basicTemplate.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_VALUE_LEN, 32));
            template.merge(basicTemplate);
        }
        CK_MECHANISM ck_mechanism = new CK_MECHANISM(cipherEnum.getKeyGenerationMechanism(), ((long) 0));
        keyHandle = pkcs11.C_GenerateKey(sessionHandle, ck_mechanism, template.finishAndGetAttributes());
        initialize();
    }

    @Override
    public byte[] encrypt(byte[] dataToEncrypt) throws PKCS11Exception {
        if (!super.isInitialized()) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
//            throw new IllegalStateException("The key is not initialized. Use generate() or create() first.");
        } else {
            CK_MECHANISM ck_mechanism = new CK_MECHANISM(cipherEnum.getCipherMechanism());
            pkcs11.C_EncryptInit(sessionHandle, ck_mechanism, keyHandle);
            //only supports one-block encrypted messages
            byte[] toEncrypt = new byte[blockSize];
            //padding
            for (int i = 0; i < toEncrypt.length; i++) {
                if (dataToEncrypt.length > i) {
                    toEncrypt[i] = dataToEncrypt[i];
                } else {
                    toEncrypt[i] = 0;
                }
            }
            byte[] encrypted = new byte[blockSize];
            pkcs11.C_Encrypt(sessionHandle, toEncrypt, 0, toEncrypt.length, encrypted, 0, encrypted.length);
            return encrypted;
        }
    }

    @Override
    public byte[] decrypt(byte[] dataToDecrypt) throws PKCS11Exception {
        CK_MECHANISM mechanism = new CK_MECHANISM(cipherEnum.getCipherMechanism());
        pkcs11.C_DecryptInit(sessionHandle, mechanism, keyHandle);

        byte[] decrypted = new byte[dataToDecrypt.length];
        pkcs11.C_Decrypt(sessionHandle, dataToDecrypt, 0, dataToDecrypt.length, decrypted, 0, decrypted.length);
        return decrypted;
    }

    @Override
    public byte[] wrapKey(Key keyToWrap) throws PKCS11Exception {
        if (!super.isInitialized()) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
//            throw new IllegalStateException("The key is not initialized. Use generate() or create() first.");
        } else {
            CK_MECHANISM ck_mechanism = cipherEnum.getCipherMechanismReady();
            byte[] wrappedKey = pkcs11.C_WrapKey(sessionHandle, ck_mechanism, keyHandle, keyToWrap.getKeyHandle());
            return wrappedKey;
        }
    }

    @Override
    public SecretKey unwrapSecretKey(byte[] keyToUnwrap, Template template, CipherEnum cipher) throws PKCS11Exception {
        if (!super.isInitialized()) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
//            throw new IllegalStateException("The key is not initialized. Use generate() or create() first.");
        } else {
            CK_MECHANISM unwrappingMechanism = new CK_MECHANISM(cipherEnum.getCipherMechanism());
            SecretKey unwrappedKey = new SecretKey(sessionHandle, pkcs11, cipher);
            if (template == null) {
                template = new Template();
                Template basicTemplate = templateMaker.getBasicSecretKeyTemplate(cipher);
                template.merge(basicTemplate);
            }
            unwrappedKey.setKeyHandleAndInitialize(pkcs11.C_UnwrapKey(sessionHandle, unwrappingMechanism, keyHandle, keyToUnwrap, template.finishAndGetAttributes()));
            return unwrappedKey;
        }
    }

    @Override
    public PrivateKey unwrapPrivateKey(byte[] toUnwrap, Template template, PublicKey correspondingPublicKey) throws PKCS11Exception {
        if (!super.isInitialized()) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
//            throw new IllegalStateException("The key is not initialized. Use generate() or create() first.");
        } else {
            CK_MECHANISM unwrappingMechanism = new CK_MECHANISM(cipherEnum.getCipherMechanism());
            PrivateKey unwrappedKey = new PrivateKey(sessionHandle, pkcs11, CipherEnum.RSA); //only asymmetric support is for RSA
            unwrappedKey.setKeyHandleAndInitialize(pkcs11.C_UnwrapKey(sessionHandle, unwrappingMechanism, keyHandle, toUnwrap, template.finishAndGetAttributes()));
            unwrappedKey.setPublicKey(correspondingPublicKey);
            return unwrappedKey;
        }
    }

    /**
     * Attempts to create the key for the given cipher. The template may be
     * specified.
     *
     * @param template The optional template for the key creation. If it is
     * null, default values are used instead.
     */
    public void create(Template template) throws PKCS11Exception {
        if (template == null) {
            Template basicTemplate = new Template();
            basicTemplate.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_SECRET_KEY));
            basicTemplate.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, cipherEnum.getKeyType()));
            byte[] randomValue = new byte[24]; //we can safely use 24-byte long keys, since both DES3 and AES permit 24 byte long keys.
            random.nextBytes(randomValue);
            basicTemplate.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_VALUE, randomValue));
            template = new Template();
            template.merge(basicTemplate);
        }
        keyHandle = pkcs11.C_CreateObject(sessionHandle, template.finishAndGetAttributes());
        initialize();
    }

    @Override
    public Key copyObject(Template template) throws PKCS11Exception {
        SecretKey copiedKey = new SecretKey(sessionHandle, pkcs11, cipherEnum);
        copiedKey.setKeyHandleAndInitialize(pkcs11.C_CopyObject(sessionHandle, keyHandle, template.finishAndGetAttributes()));
        return copiedKey;
    }

    @Override
    public Key deriveKey(CK_MECHANISM mechanism, Template template) throws PKCS11Exception {
        SecretKey derivedKey = new SecretKey(sessionHandle, pkcs11, cipherEnum);
        derivedKey.setKeyHandleAndInitialize(pkcs11.C_DeriveKey(sessionHandle, mechanism, keyHandle, template.finishAndGetAttributes()));
        return derivedKey;
    }
}
